# -*- coding: utf-8 -*-
# @Author: Tim Liu
# @Date: 2024-07-07
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-07-07

# @Description: Audio Document Loader for RAG

from typing import AsyncIterator, Iterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document

from openai import AzureOpenAI

from config.settings import *

from crewplus.helper.fileutil import FileUtil

class AudioDocumentLoader(BaseLoader):
    """An audio document loader that loads an audio file into doc."""

    def __init__(self, file_path: str) -> None:
        """Initialize the loader with a file path.

        Args:
            file_path: The path to the file to load.
        """
        self.file_path = file_path

    def lazy_load(self) -> Iterator[Document]:  # <-- Does not take any arguments
        """A lazy loader that reads a file line by line.

        When you're implementing lazy load methods, you should use a generator
        to yield documents one by one.
        """
        audio_data= FileUtil.download_file(self.file_path)

        client = AzureOpenAI(
            api_key=WHISPER_AZURE_OPENAI_API_KEY,
            api_version=WHISPER_AZURE_OPENAI_API_VERSION,
            azure_endpoint = WHISPER_AZURE_OPENAI_ENDPOINT
        )

        data = {
          'file': (f'audio_file.wav', audio_data, 'audio/wav')
        }
        
        result = client.audio.transcriptions.create(
            file=data['file'],
            model=WHISPER_AZURE_OPENAI_CHAT_DEPLOYMENT_NAME
        )

        yield Document(
            page_content=result.text,
            metadata={"file_type": "audio", "source": self.file_path},
        )